import torch
import torch.nn as nn
from torchvision.models import resnet18, resnet34, resnet50, vgg16, vgg11, alexnet, densenet121, resnet101
from timm.models import vit_small_patch16_224, vit_base_patch16_224, vit_large_patch16_224, vit_base_patch32_224, \
    vit_large_patch32_224, vit_base_resnet26d_224, vit_base_resnet50d_224, vit_small_resnet26d_224


# 接触角回归预测模型
class ContactAngle(nn.Module):
    def __init__(self, model_name='Resnet18', out_puts=1020):
        super(ContactAngle, self).__init__()
        # 图像特征提取网络
        self.model_name = model_name
        # 图像特征输出向量维度
        self.out_puts = out_puts
        self.model = Select_model(name=self.model_name, pretrained=False, out=self.out_puts)
        # 图像特征与加工参数融合的多层感知机层
        self.fc = nn.Sequential(
            # (N, 1024) -> (N, 512)
            nn.Linear(in_features=self.out_puts + 4, out_features=512),
            nn.ReLU(inplace=True),
            # (N, 512) -> (N, 256)
            nn.Linear(in_features=512, out_features=256),
            nn.ReLU(inplace=True),
            # (N, 256) -> (N, 1)
            nn.Linear(in_features=256, out_features=1)
        )

    def forward(self, x1, x2):
        # (N, 3, 224, 224) -> (N, 1020)
        out1 = self.model(x1)
        # (N, 1020) + (N, 4) -> (N, 1024)
        out2 = torch.cat((out1, x2), dim=-1)
        # (N, 1024) -> (N, 1)
        out = self.fc(out2)
        return out


def Select_model(name='Resnet18', pretrained=False, out=2):
    pre_trained = pretrained
    out_num = out
    paras = {
        'Densnet121': Densnet121(pretrained=pre_trained, out=out_num),
        'Resnet18': Resnet18(pretrained=pre_trained, out=out_num),
        'Resnet34': Resnet34(pretrained=pre_trained, out=out_num),
        'Resnet50': Resnet50(pretrained=pre_trained, out=out_num),
        'Resnet101': Resnet50(pretrained=pre_trained, out=out_num),
        'Vgg16': Vgg16(pretrained=pre_trained, out=out_num),
        'Vgg11': Vgg11(pretrained=pre_trained, out=out_num),
        'Alexnet': Alexnet(pretrained=pre_trained, out=out_num),
        'Vit_small_patch16_224': Vit_small_patch16_224(pretrained=pre_trained, out=out_num),
        'Vit_base_patch16_224': Vit_base_patch16_224(pretrained=pre_trained, out=out_num),
        'Vit_base_patch32_224': Vit_base_patch32_224(pretrained=pre_trained, out=out_num),
        'Vit_large_patch16_224': Vit_large_patch16_224(pretrained=pre_trained, out=out_num),
        'Vit_large_patch32_224': Vit_large_patch32_224(pretrained=pre_trained, out=out_num),
        'Vit_base_resnet26d_224': Vit_base_resnet26d_224(pretrained=pre_trained, out=out_num),
        'Vit_base_resnet50d_224': Vit_base_resnet50d_224(pretrained=pre_trained, out=out_num),
        'Vit_small_resnet26d_224': Vit_small_resnet26d_224(pretrained=pre_trained, out=out_num)
    }
    return paras[name]


def Densnet121(pretrained=False, out=1020):
    model_out = densenet121(pretrained=pretrained)
    model_out.classifier = nn.Linear(in_features=1024, out_features=out, bias=True)
    return model_out


def Resnet18(pretrained=False, out=1020):
    model_out = resnet18(pretrained=pretrained)
    model_out.fc = nn.Linear(in_features=512, out_features=out, bias=True)
    return model_out


def Resnet34(pretrained=False, out=1020):
    model_out = resnet34(pretrained=pretrained)
    model_out.fc = nn.Linear(in_features=512, out_features=out, bias=True)
    return model_out


def Resnet50(pretrained=False, out=1020):
    model_out = resnet50(pretrained=pretrained)
    model_out.fc = nn.Linear(in_features=2048, out_features=out, bias=True)
    return model_out

def Resnet101(pretrained=False, out=1020):
    model_out = resnet101(pretrained=pretrained)
    model_out.fc = nn.Linear(in_features=2048, out_features=out, bias=True)
    return model_out


def Vgg16(pretrained=False, out=1020):
    model_out = vgg16(pretrained=pretrained)
    model_out.classifier[6] = nn.Linear(in_features=4096, out_features=out, bias=True)
    return model_out


def Vgg11(pretrained=False, out=1020):
    model_out = vgg11(pretrained=pretrained)
    model_out.classifier[6] = nn.Linear(in_features=4096, out_features=out, bias=True)
    return model_out


def Alexnet(pretrained=False, out=1020):
    model_out = alexnet(pretrained=pretrained)
    model_out.classifier[6] = nn.Linear(in_features=4096, out_features=out, bias=True)
    return model_out


def Vit_small_patch16_224(pretrained=False, out=1020):
    model_out = vit_small_patch16_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=768, out_features=out, bias=True)
    return model_out


def Vit_base_patch16_224(pretrained=False, out=1020):
    model_out = vit_base_patch16_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=768, out_features=out, bias=True)
    return model_out


def Vit_base_patch32_224(pretrained=False, out=1020):
    model_out = vit_base_patch32_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=768, out_features=out, bias=True)
    return model_out


def Vit_large_patch16_224(pretrained=False, out=1020):
    model_out = vit_large_patch16_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=1024, out_features=out, bias=True)
    return model_out


def Vit_large_patch32_224(pretrained=False, out=1020):
    model_out = vit_large_patch32_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=1024, out_features=out, bias=True)
    return model_out


def Vit_base_resnet26d_224(pretrained=False, out=1020):
    model_out = vit_base_resnet26d_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=768, out_features=out, bias=True)
    return model_out


def Vit_base_resnet50d_224(pretrained=False, out=1020):
    model_out = vit_base_resnet50d_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=768, out_features=out, bias=True)
    return model_out


def Vit_small_resnet26d_224(pretrained=False, out=1020):
    model_out = vit_small_resnet26d_224(pretrained=pretrained)
    model_out.head = nn.Linear(in_features=768, out_features=out, bias=True)
    return model_out


if __name__ == '__main__':
    model = ContactAngle(out_puts=1020)
